#include <stdio.h>
#include <stdbool.h>
#include <stdint.h>
#include <assert.h>
#include <omp.h>
#include <immintrin.h>
#include <chrono>
#include <iostream>
#include <thread>
void Cumsum_cx1(const float *input, float *output, 
                  int out_dim, int axis_dim, int inner_dim, 
                  bool exclusive) {
    int i = 0, j = 0, k = 0;
    
    // 初始化阶段（仅调整指针起始位置）
    if (!exclusive) {
        for (i = 0; i < out_dim; ++i) {
            const float *layer_input = input + i * axis_dim * inner_dim;  // 改为正向起始位置
            float *layer_output = output + i * axis_dim * inner_dim;      // 改为正向起始位置
            for (j = 0; j < inner_dim; ++j) {
                *(layer_output + j) = *(layer_input + j);  // 保持复制逻辑
            }
        }
    } else {
        for (i = 0; i < out_dim; ++i) {
            float *layer_output = output + i * axis_dim * inner_dim;  // 改为正向起始位置
            for (j = 0; j < inner_dim; ++j) {
                *(layer_output + j) = 0.0f;  // 保持初始化逻辑
            }
        }
    }

    // 累加阶段（调整指针方向为正向）
    int input_offset = exclusive ? 0 : 1;
    for (i = 0; i < out_dim; ++i) {
        const float *layer_input = input + i * axis_dim * inner_dim + input_offset * inner_dim;  // 正向偏移
        float *layer_last_output = output + i * axis_dim * inner_dim;                           // 正向起始
        float *layer_output = layer_last_output + inner_dim;                                     // 正向步进

        for (j = 1; j < axis_dim; ++j) {
            for (k = 0; k < inner_dim; ++k) {
                *(layer_output + k) = *(layer_input + k) + *(layer_last_output + k);  // 改为正向指针
            }
            layer_input += inner_dim;      // 正向移动
            layer_last_output += inner_dim;
            layer_output += inner_dim;
        }
    }
}


void Cumsum_cx(const float *input, float *output,
                     int out_dim, int axis_dim, int inner_dim,
                     bool exclusive) {
    const int slice_sz = axis_dim * inner_dim;

    
    for (int i = 0; i < out_dim; ++i) {
        const float *in_base  = input  + i * slice_sz;
        float       *out_base = output + i * slice_sz;

        // — 初始化第 0 行 —
        if (exclusive) {
            for (int k = 0; k < inner_dim; ++k) {
                out_base[k] = 0.0f;
            }
        } else {
            for (int k = 0; k < inner_dim; ++k) {
                out_base[k] = in_base[k];
            }
        }

        // — 依次累加第 1 … axis_dim-1 行 —
        for (int j = 1; j < axis_dim; ++j) {
            // 计算本行输入指针：exclusive 时读 j-1，否则读 j
            const float *cur_in   = in_base
                                  + (exclusive ? (j - 1) : j) * inner_dim;
            float       *prev_out = out_base + (j - 1) * inner_dim;
            float       *cur_out  = out_base + j * inner_dim;

            int k = 0;
            // AVX2 每次处理 8 个 float
            for (; k + 8 <= inner_dim; k += 8) {
                __m256 vin  = _mm256_loadu_ps(cur_in   + k);
                __m256 vpre = _mm256_loadu_ps(prev_out + k);
                _mm256_storeu_ps(cur_out + k,
                                 _mm256_add_ps(vin, vpre));
            }
            // 处理剩余不足 8 的部分
            for (; k < inner_dim; ++k) {
                cur_out[k] = cur_in[k] + prev_out[k];
            }
        }
    }
}





//void Cumsum_Hybrid(const float *input, float *output,
void Cumsum_Hybrid(const float *input, float *output,
                     int out_dim, int axis_dim, int inner_dim,
                     bool exclusive) {
    const int slice_sz = axis_dim * inner_dim;

    #pragma omp parallel for num_threads(3)
    for (int i = 0; i < out_dim; ++i) {
        const float *in_base  = input  + i * slice_sz;
        float       *out_base = output + i * slice_sz;

        // — 初始化第 0 行 —
        if (exclusive) {
            for (int k = 0; k < inner_dim; ++k) {
                out_base[k] = 0.0f;
            }
        } else {
            for (int k = 0; k < inner_dim; ++k) {
                out_base[k] = in_base[k];
            }
        }

        // — 依次累加第 1 … axis_dim-1 行 —
        for (int j = 1; j < axis_dim; ++j) {
            // 计算本行输入指针：exclusive 时读 j-1，否则读 j
            const float *cur_in   = in_base
                                  + (exclusive ? (j - 1) : j) * inner_dim;
            float       *prev_out = out_base + (j - 1) * inner_dim;
            float       *cur_out  = out_base + j * inner_dim;

            int k = 0;
            // AVX2 每次处理 8 个 float
            for (; k + 8 <= inner_dim; k += 8) {
                __m256 vin  = _mm256_loadu_ps(cur_in   + k);
                __m256 vpre = _mm256_loadu_ps(prev_out + k);
                _mm256_storeu_ps(cur_out + k,
                                 _mm256_add_ps(vin, vpre));
            }
            // 处理剩余不足 8 的部分
            for (; k < inner_dim; ++k) {
                cur_out[k] = cur_in[k] + prev_out[k];
            }
        }
    }
}

void Cumsum_Hybrid1(const float *input, float *output,
                   int out_dim, int axis_dim, int inner_dim,
                   bool exclusive) {
    const int slice_sz = axis_dim * inner_dim;

    // 允许 OpenMP 根据负载自动调整线程数（可选）
    omp_set_dynamic(4);

    // 并行 for，3 个线程，动态调度，每次分配 1 个 i
    #pragma omp parallel for num_threads(1) schedule(dynamic, 1)
    for (int i = 0; i < out_dim; ++i) {
        const float *in_base  = input  + i * slice_sz;
        float       *out_base = output + i * slice_sz;

        // 初始化第 0 行
        if (exclusive) {
            int k = 0;
            for (; k + 8 <= inner_dim; k += 8) {
                _mm256_storeu_ps(out_base + k, _mm256_setzero_ps());
            }
            for (; k < inner_dim; ++k) {
                out_base[k] = 0.0f;
            }
        } else {
            int k = 0;
            for (; k + 8 <= inner_dim; k += 8) {
                __m256 vin = _mm256_loadu_ps(in_base + k);
                _mm256_storeu_ps(out_base + k, vin);
            }
            for (; k < inner_dim; ++k) {
                out_base[k] = in_base[k];
            }
        }

        // 依次累加第 1 … axis_dim-1 行
        for (int j = 1; j < axis_dim; ++j) {
            const float *cur_in   = in_base
                                  + (exclusive ? (j - 1) : j) * inner_dim;
            float       *prev_out = out_base + (j - 1) * inner_dim;
            float       *cur_out  = out_base + j * inner_dim;

            int k = 0;
            for (; k + 8 <= inner_dim; k += 8) {
                __m256 vin  = _mm256_loadu_ps(cur_in   + k);
                __m256 vpre = _mm256_loadu_ps(prev_out + k);
                __m256 vsum = _mm256_add_ps(vin, vpre);
                _mm256_storeu_ps(cur_out + k, vsum);
            }
            for (; k < inner_dim; ++k) {
                cur_out[k] = cur_in[k] + prev_out[k];
            }
        }
    }
}
void test_CumsumHybrid() {
    float input[72000] ;
    float output[72000];
    int out_dim = 12000, axis_dim = 3, inner_dim = 2;
    for(int i=0;i<72000;i++){
        input[i]=(float)1;
    }

    // 测试 exclusive 模式
   
    

   clock_t start = clock();
    for(int i=0;i<1000;i++){ Cumsum_Hybrid(input, output, out_dim, axis_dim, inner_dim, true);}
   clock_t end = clock();
    double duration = (double)(end - start) / CLOCKS_PER_SEC;
    printf("st parallel time:\n");
     double time = duration ;
      printf("%f ", time);

  start = clock();
  
     for(int i=0;i<1000;i++){ Cumsum_Hybrid1(input, output, out_dim, axis_dim, inner_dim, true);}
    end = clock();
   duration = (double)(end - start) / CLOCKS_PER_SEC;
    printf("dy time:\n");
      time = duration ;
      printf("%f ", time);




    // 打印结果
   printf("Output (exclusive=true):\n");
    for (int i = 0; i < 120; ++i) {
        printf("%.1f ", output[i]);
        if ((i + 1) % (axis_dim * inner_dim) == 0) printf("\n");
    }
}




int main() {

    test_CumsumHybrid();
 
    return 0;
}